#!/usr/bin/python3

import base64
import json
import logging
import os
import shutil
import sys
import uuid


def send_control(**msg):
    logging.debug(msg)
    msg = f'\n{json.dumps(msg)}\n'.encode('utf-8')
    # yes, we're writing to stdin: it's a socket, and for askpass, ssh captures stdout
    os.write(0, f'{len(msg)}\n'.encode('utf-8') + msg)


def recv_control():
    # using buffered IO is safe here: we only get one message sent at a
    # time, and then it's our turn to say something before the next one
    # will be sent.  therefore, it's impossible to over-read.
    fp = sys.stdin.buffer
    length = int(fp.readline())
    # take advantage of empty channel name being json whitespace
    return json.loads(fp.read(length))


def do_authorize(challenge='*', **kwargs):
    cookie = str(uuid.uuid4())
    send_control(command='authorize', challenge=challenge, cookie=cookie, **kwargs)
    response = recv_control()
    assert response['cookie'] == cookie
    logging.debug(response)
    if 'response' not in response:
        raise ValueError('invalid authorize response')
    return response['response'].split(' ')


def base64_encode(s):
    return base64.b64encode(s.encode('utf-8')).decode('utf-8')


def base64_decode(s):
    return base64.b64decode(s.encode('utf-8')).decode('utf-8')


def cockpit_client_ssh():
    # this is pretty evil, but oh well...
    xdg_data_home = os.getenv('XDG_DATA_HOME') or os.path.expanduser('~/.local/share')
    askpass = os.path.join(xdg_data_home, 'cockpit-client-askpass')
    shutil.copy(__file__, askpass, follow_symlinks=False)

    connect_to = sys.argv[1]
    if connect_to == 'localhost':
        cmd = ['flatpak-spawn', '--host', 'cockpit-bridge']
    else:
        host, _, port = connect_to.rpartition(':')
        # catch cases like `host:123` but not cases like `[2001:abcd::1]`
        if port.isdigit():
            host_args = ['-p', port, host]
        else:
            host_args = [connect_to]

        cmd = ['flatpak-spawn', '--host', '/usr/bin/env',
               'SSH_ASKPASS_REQUIRE=force', f'SSH_ASKPASS={askpass}',
               '/usr/bin/ssh'] + host_args + ['cockpit-bridge']

    os.execvp(cmd[0], cmd)


def cockpit_client_askpass():
    message, _, prompt = sys.argv[1].rpartition('\n')
    challenge = f'X-Conversation - {base64_encode(prompt)}'
    response = do_authorize(challenge, message=message)
    if len(response) == 3:  # will be 2 if empty
        output = base64_decode(response[2])
        logging.debug(output)
        print(output)


if __name__ == '__main__':
    # logging.basicConfig(level=logging.DEBUG)

    logging.debug(sys.argv)
    if 'cockpit-client-ssh' in sys.argv[0]:
        cockpit_client_ssh()
    elif 'cockpit-client-askpass' in sys.argv[0]:
        cockpit_client_askpass()
    else:
        raise Exception(f'unknown executable name {sys.argv[0]}')
